home *** CD-ROM | disk | FTP | other *** search
/ MacHack 2001 / MacHack 2001.toast / pc / The Hacks / Unprotected Memory / MyProcessWorker.m < prev    next >
Encoding:
Text File  |  2001-06-23  |  5.2 KB  |  214 lines
  1. #import "MyProcessWorker.h"
  2.  
  3. #include <mach/mach_types.h>
  4. #include <mach/mach_traps.h>
  5. #include <mach/mach_error.h>
  6.  
  7. #include <sys/sysctl.h>
  8. #include <sys/reboot.h>
  9. #include <unistd.h>
  10.  
  11. @interface MyProcessWorker(PrivateMethods)
  12.  
  13. - (void)notifyCrash;
  14.  
  15. @end
  16.  
  17. @implementation MyProcessWorker
  18.  
  19. - (id)init
  20. {
  21.     if (self = [super init])
  22.     {
  23.         [self walkProcessList];
  24.     }
  25.     
  26.     return self;
  27. }
  28.  
  29. - (void)walkProcessList
  30. {
  31.     int                    mib[4];
  32.     int                 i;
  33.     size_t                 bufSize;
  34.     int                 nentries;
  35.     struct kinfo_proc     *kprocbuf = NULL, *kp;
  36.     BOOL                initialWalk = (myLastSeenPID == 0);
  37.     
  38.     NS_DURING
  39.     {
  40.         mib[0] = CTL_KERN;
  41.         mib[1] = KERN_PROC;
  42.         mib[2] = KERN_PROC_ALL;
  43.         mib[3] = 0;
  44.     
  45.         if (sysctl(mib, 4, NULL, &bufSize, NULL, 0) < 0)
  46.         {
  47.             [NSException raise:NSInternalInconsistencyException
  48.                         format:@"unable to count processes (sysctl)"];
  49.         }
  50.         
  51.         nentries = bufSize / sizeof(struct kinfo_proc);
  52.         kprocbuf = kp = (struct kinfo_proc *)NSZoneMalloc(NULL, bufSize);
  53.  
  54.         if (sysctl(mib, 4, kp, &bufSize, NULL, 0) < 0)
  55.         {
  56.             [NSException raise:NSInternalInconsistencyException
  57.                         format:@"unable to get process list (sysctl)"];
  58.         }   
  59.         
  60.         // it lies the first time, fortunately with a buffer that is too big
  61.         nentries = bufSize / sizeof(struct kinfo_proc);
  62.     
  63.         if (kp != NULL)
  64.         {
  65.             for(i=nentries; --i>=0; ++kp)
  66.             {
  67.                 if (kp->kp_proc.p_pid > myLastSeenPID)
  68.                 {
  69.                     myLastSeenPID = kp->kp_proc.p_pid;
  70.  
  71.                     if ( ! initialWalk)
  72.                     {
  73.                         [NSThread detachNewThreadSelector:@selector(myThreadMain:)
  74.                             toTarget:self withObject:[NSNumber numberWithInt:myLastSeenPID]];
  75.                     }
  76.                 }
  77.             }
  78.         }
  79.     }
  80.     NS_HANDLER
  81.     {
  82.         NSLog(@"Couldn't install exception handler for pid: %d; Exception :%@", [localException description]);
  83.     }
  84.     NS_ENDHANDLER;
  85.     
  86.     if (kprocbuf != NULL)
  87.         NSZoneFree(NULL, kprocbuf);
  88. }
  89.  
  90. #define MACH_CHECK_ERROR(name,ret)    \
  91.         if (ret != KERN_SUCCESS) {    \
  92.         mach_error(#name, ret);        \
  93.         exit(1);                    \
  94. } 
  95.  
  96. - (void)myThreadMain:(id)arg
  97. {
  98.     const int             MSG_SIZE = 512;
  99.     id                        pid = [arg retain];
  100.     mach_msg_header_t     *msg, *reply;
  101.     kern_return_t         krc;
  102.     mach_port_t         _exceptionPort;
  103.     task_t                 targetTask;
  104.     
  105.     NS_DURING
  106.     {
  107.         krc = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &_exceptionPort);
  108.         if (krc != KERN_SUCCESS)
  109.         {
  110.             [NSException raise: NSInternalInconsistencyException format: @"Unable to create handler port, krc = %d, %s", krc, mach_error_string(krc)];
  111.         }
  112.         
  113.         krc = mach_port_insert_right(mach_task_self(), _exceptionPort, _exceptionPort, MACH_MSG_TYPE_MAKE_SEND);
  114.         if (krc != KERN_SUCCESS)
  115.         {
  116.             [NSException raise: NSInternalInconsistencyException format: @"Unable insert send write into handler port, krc = %d, %s", krc, mach_error_string(krc)];
  117.         }
  118.     
  119.         krc = task_for_pid(mach_task_self(), [pid intValue], &targetTask);
  120.         if (krc != KERN_SUCCESS)
  121.         {
  122.             [NSException raise: NSInternalInconsistencyException format: @"task_for_pid failed, krc = %d, %s", krc, mach_error_string(krc)];
  123.         }
  124.         
  125.         krc = task_set_exception_ports(targetTask, EXC_MASK_BAD_ACCESS | EXC_MASK_BAD_INSTRUCTION | EXC_MASK_ARITHMETIC, _exceptionPort, EXCEPTION_DEFAULT, THREAD_STATE_NONE);
  126.         if (krc != KERN_SUCCESS)
  127.         {
  128.             [NSException raise: NSInternalInconsistencyException format: @"task_set_exception_ports failed, krc = %d, %s", krc, mach_error_string(krc)];
  129.         }
  130.     
  131.         msg = alloca(MSG_SIZE);
  132.         reply = alloca(MSG_SIZE);
  133.  
  134.         while (TRUE)
  135.         {
  136.             NSAutoreleasePool    *pool = [[NSAutoreleasePool alloc] init];
  137.             
  138.             // we only get what we consider to be fatal messages,
  139.             // so just shoot the process and do our thing
  140.  
  141.             krc = mach_msg(msg, MACH_RCV_MSG, MSG_SIZE, MSG_SIZE, _exceptionPort, 0, MACH_PORT_NULL);
  142.             if (krc == KERN_SUCCESS)
  143.             {
  144.                 kill([pid intValue], SIGKILL);
  145.                 [self notifyCrash];
  146.             }
  147.  
  148.             [pool release];
  149.         }
  150.     }
  151.     NS_HANDLER
  152.     {
  153.     
  154.     }
  155.     NS_ENDHANDLER;
  156. }
  157.  
  158. @end
  159.  
  160. @implementation MyProcessWorker(PrivateMethods)
  161.  
  162. - (void)notifyCrash
  163. {
  164.     const int         FREE_CRASH_LIMIT = 1;
  165.     CFStringRef     APPLICATION_PREF_ID = CFSTR("Unprotected Memory");
  166.     CFStringRef     ALREADY_REBOOTED_KEY = CFSTR("AlreadyRebooted");
  167.     
  168.     BOOL            noGrace = NO;
  169.     CFTimeInterval    timeout = 5;
  170.  
  171.     myCrashCount++;
  172.     
  173.     if ( CFPreferencesGetAppBooleanValue(ALREADY_REBOOTED_KEY, APPLICATION_PREF_ID, NULL) )
  174.     {
  175.         noGrace = YES;
  176.         timeout = 0.5;
  177.     }
  178.     
  179.     if (noGrace || (myCrashCount > FREE_CRASH_LIMIT))
  180.     {
  181.         SInt32            response;
  182.         CFOptionFlags    responseFlags;
  183.     
  184.         response = CFUserNotificationDisplayAlert(timeout, 0L, NULL, NULL, NULL,
  185.                                                     CFSTR("An application has unexpectedly quit."),
  186.                                                     CFSTR("You have 5 seconds to save all your work before the machine reboots."),
  187.                                                     CFSTR("Restart"),
  188.                                                     NULL,
  189.                                                     NULL,
  190.                                                     &responseFlags);
  191.                                                     
  192.         CFPreferencesSetAppValue(ALREADY_REBOOTED_KEY, kCFBooleanTrue, APPLICATION_PREF_ID);
  193.         CFPreferencesAppSynchronize(APPLICATION_PREF_ID);
  194.  
  195.                                                     
  196.         if (0 != reboot(RB_AUTOBOOT))
  197.         {
  198.             // call only returns if it could fail, otherwise
  199.             // the machine is rebooting
  200.             
  201.             NSLog(@"reboot didn't work, did you run it as root dummy?");
  202.         }
  203.     }
  204.     else
  205.     {
  206.         (void)CFUserNotificationDisplayNotice(0, 0L, NULL, NULL, NULL,
  207.                                               CFSTR("An application has unexpectedly quit."),
  208.                                               CFSTR("Aren't you glad you didn't have to reboot your machine?"),
  209.                                               CFSTR("Whew!"));
  210.     }
  211. }
  212.  
  213. @end
  214.